package com.wind.common.utils;

import cn.hutool.core.thread.NamedThreadFactory;
import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.*;

/**
 * 线程池工具
 *
 * @author FengYin
 **/
@Slf4j
public class ThreadPoolUtil {

    private static final ExecutorService executorService;

    static {
        int corePoolSize = Runtime.getRuntime().availableProcessors() * 2;
        log.info("初始化主线程池，核心线程数:{}", corePoolSize);
        executorService = new ThreadPoolExecutor(corePoolSize,
                corePoolSize,
                60,
                TimeUnit.SECONDS,
                new LinkedBlockingQueue<>(),
                new NamedThreadFactory("pool-", true));
    }

    /**
     * 获取线程池
     *
     * @return 线程池
     */
    private static ExecutorService executorService() {
        return executorService;
    }

    /**
     * 执行线程
     *
     * @param command 线程
     */
    public static void execute(Runnable command) {
        execute(null, command);
    }

    /**
     * 执行线程
     *
     * @param executorService 线程池，为空则使用系统默认线程池
     * @param command         线程
     */

    public static void execute(ExecutorService executorService, Runnable command) {
        if (executorService == null) {
            executorService = executorService();
        }
        executorService.execute(wrap(command));
    }

    /**
     * 提交任务（有返回）
     *
     * @param task 异步任务
     * @param <T>  返回类型
     * @return 任务返回
     */
    public static <T> Future<T> submit(Callable<T> task) {
        return submit(null, task);
    }

    /**
     * 提交任务 (有返回)
     *
     * @param executorService 线程池，为空则使用系统默认线程池
     * @param task            异步任务
     * @param <T>             返回类型
     * @return 任务返回
     */
    public static <T> Future<T> submit(ExecutorService executorService, Callable<T> task) {
        if (executorService == null) {
            executorService = executorService();
        }
        return executorService.submit(wrap(task));
    }


    /**
     * 提交任务（无返回）
     *
     * @param task 异步任务
     * @return 任务返回
     */
    public static Future<Void> submit(Runnable task) {
        return submit(null, task);
    }

    /**
     * 提交任务 (无返回)
     *
     * @param executorService 线程池，为空则使用系统默认线程池
     * @param task            异步任务
     * @return 任务返回
     */
    public static Future<Void> submit(ExecutorService executorService, Runnable task) {
        return submit(executorService, () -> {
            task.run();
            return null;
        });
    }


    /**
     * 为Callable任务添加TraceId传递和异常捕获的装饰器方法.
     * <p>
     * 该方法用于装饰Callable任务,确保在任务执行前后能正确地传递和清理TraceId信息,
     * 并且在任务执行过程中捕获并转换为运行时异常的异常.
     *
     * @param <T>      泛型参数，表示Callable任务的返回类型
     * @param callable 需要被装饰的Callable任务
     * @return 装饰后的Callable任务，它包含了TraceId的传递和异常捕获逻辑
     */
    public static <T> Callable<T> wrap(final Callable<T> callable) {
        // 获取当前线程的TraceId
        String traceId = TraceLogContext.getTraceId();
        // 返回一个新的Callable，它包含了TraceId传递和异常捕获的逻辑
        return () -> {
            try {
                // 在执行任务前设置TraceId到当前线程
                TraceLogContext.setTraceId(traceId);
                // 执行原始的Callable任务并返回结果
                return callable.call();
            } catch (Exception e) {
                // 捕获Callable任务执行过程中抛出的异常，并记录错误日志
                log.error("线程执行异常", e);
                // 将捕获的异常转换为运行时异常并抛出
                throw new RuntimeException(e);
            } finally {
                // 无论任务执行是否成功，最终清除当前线程的TraceId
                TraceLogContext.removeTraceId();
            }
        };
    }


    /**
     * 包装一个Runnable任务，用于自动管理TraceId
     * 在任务执行前后，自动设置和清除TraceId，以便于链路追踪
     * 如果任务执行过程中抛出异常，会捕获并记录该异常
     *
     * @param runnable 被包装的Runnable任务
     * @return 包装后的Runnable任务
     */
    public static Runnable wrap(final Runnable runnable) {
        // 获取当前线程的TraceId，以便在后续的执行中保持一致性
        String traceId = TraceLogContext.getTraceId();
        // 返回一个新的Runnable任务，该任务在执行时会自动管理TraceId
        return () -> {
            try {
                // 在任务执行前设置当前线程的TraceId
                TraceLogContext.setTraceId(traceId);
                // 执行原始的Runnable任务
                runnable.run();
            } catch (Exception e) {
                // 如果任务执行过程中抛出异常，记录该异常
                log.error("线程执行异常", e);
            } finally {
                // 无论任务执行是否成功，最终清除当前线程的TraceId
                TraceLogContext.removeTraceId();
            }
        };
    }

}
